from .module import PREPROCESS
import cv2
import numpy as np


@PREPROCESS.register_module
class Resize:

    def __init__(self, size, random_crop):
        self.size = size
        self.pre_size = (256, 256)
        self.random_crop = random_crop

    def process(self, img, anno):
        anno = anno.copy()
        h, w, c = img.shape
        if self.random_crop:
            img = cv2.resize(img, dsize=self.pre_size)
            r_x = np.random.randint(self.pre_size[1] - self.size[1])
            r_y = np.random.randint(self.pre_size[0] - self.size[0])
            img = img[r_y:r_y + self.size[0], r_x:r_x + self.size[1]]
            # anno is not correct here, but it's ok for training
        else:
            img = cv2.resize(img, dsize=tuple(self.size))
            anno[1] *= self.size[0] / w
            anno[2] *= self.size[1] / h
            anno[3] *= self.size[0] / w
            anno[4] *= self.size[1] / h

        # OLD VERSION OF RANDOM CROP
        # if self.random_crop:
        #     temp = np.zeros(shape=img.shape, dtype=img.dtype)
        #     r_x = np.random.randint(16) - 8
        #     r_y = np.random.randint(16) - 8
        #     anno[1] += r_x
        #     anno[2] += r_y
        #     anno[3] += r_x
        #     anno[4] += r_y
        #     if r_x >= 0:
        #         if r_y >= 0:
        #             temp[r_y:, r_x:] = img[0:self.size[1]-r_y, 0:self.size[0]-r_x]
        #         else:
        #             temp[0:self.size[1]-abs(r_y), r_x:] = img[abs(r_y):, 0:self.size[0]-r_x]
        #     else:
        #         if r_y >= 0:
        #             temp[r_y:, 0:self.size[0]-abs(r_x)] = img[0:self.size[1]-r_y, abs(r_x):]
        #         else:
        #             temp[0:self.size[1]-abs(r_y), 0:self.size[0]-abs(r_x)] = img[abs(r_y):, abs(r_x):]
        #     img = temp

        return img, anno